import torch
from networks.ViT_Enhanced import ViTEnhancedNet

def load_model(checkpoint_path):
    # 设置设备
    device = torch.device("cuda:3" if torch.cuda.is_available() else "cpu")
    
    # 创建模型实例
    model = ViTEnhancedNet(num_class=8, num_head=2, pretrained=True)
    model.to(device)
    
    # 加载checkpoint
    checkpoint = torch.load(checkpoint_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    
    # 设置为评估模式
    model.eval()
    
    print(f"成功加载模型，准确率: {checkpoint_path.split('_acc')[1].split('_')[0]}")
    return model

if __name__ == "__main__":
    checkpoint_path = "/data/yxy/erp情绪识别/checkpoints/ferPlus_epoch64_acc0.7278_bacc0.7115.pth"
    model = load_model(checkpoint_path) 